import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 参数设置
a = 1.0          # 波速
L = 1.0           # 空间域长度
T = 0.5           # 时间域长度
Nx = 100         # 空间网格点数
Nt = 400         # 时间网格点数
dx = L / (Nx - 1)  # 空间步长
dt = T / (Nt - 1)  # 时间步长
c = a * dt / dx     # 波速与步长的比例

# 初始化网格
x = np.linspace(0, L, Nx)
t = np.linspace(0, T, Nt)

# 设置初始条件
u = np.sin(2 * np.pi * x)
u[0]=u[-2]
u[-1]=u[1]

# 创建图形和轴
fig, ax = plt.subplots()
line, = ax.plot(x, u, lw=2)  # 线对象用于更新动画
ax.set_xlim([0, L])
ax.set_ylim(-1, 1)  # 根据初始条件设置y轴范围
ax.set_title('Two-step Lax-Wendroff Scheme with Sin Initial Condition')

# 初始化动画
def init():
    line.set_data(x, u)
    return line,

# 更新动画的函数
def update(frame):
    # 存储前一个时间步的解
    u_prev = u.copy()
    
    # 第一步：计算半步解
    u_half = u.copy()
    for j in range(1, Nx):
        u_half[j] = (u[j] + u_prev[j - 1]) / 2 - (c / 2) * (u[j] - u_prev[j - 1])   #j-1/2
    
    # 第二步：计算全步解
    for j in range(1, Nx-1):
        u[j] = u_prev[j] - c * (u_half[j + 1] -u_half[j] ) / dx     #j+1/2,j-1/2
    
    u[0]=u[-2]
    u[-1]=u[1]
    # 更新线对象的数据
    line.set_data(x, u)
    return line,

# 创建动画
ani = FuncAnimation(fig, update, frames=Nt, init_func=init, blit=True)

# 显示动画
plt.show()